from Crypto.Cipher import AES
from Crypto.Hash import SHA256, SHA1
from Crypto.Random import get_random_bytes
import hashlib
import time
import pickle

'''
    《修改后的协议》
    1. 异或需要用int类型的值进行，且异或后的结果值的类型为int
    2. sha函数生成的值均为字节类型，若要参与计算，需要将该值设置为int，可以使用  int.from_bytes(x.digest(), byteorder='big')来转换
       sha函数的输入需要为byte，所以需要将其他类型换算为byte，在此项目中，需要将异或后的值（int）转换为byte进行计算，
       可以使用 hwA.to_bytes((hwA.bit_length() + 7) // 8, byteorder='big')  进行转换
       其他数值，byte转int可以用 int.from_bytes(message, byteorder='big')  进行转换
       ----1 byte = 8 bit
    3. get_random_bytes(1)，是随机生成1个字节的值，类型为byte
    4. 可写函数进行int和byte转换，不至于代码混乱
'''

# 全局变量
g = 3
idA = b"User_A_8150"
idB = b"User_B_6623"
pw = b"ssx2016360."


# int to byte

def intToByte(message):
    result = message.to_bytes((message.bit_length() + 7) // 8, byteorder='big')
    return result


# byte to int
def byteToInt(message):
    result = int.from_bytes(message, byteorder='big')
    return result


# 哈希函数
# def H(message):
#     return int.from_bytes(hashlib.sha1(message).digest(), byteorder='big')  # 将sha1生成的字节值转换为int

def H(data):
    # 检查输入的数据类型是否为整数
    if isinstance(data, int):
        # 如果是整数，将其转换为字节数组
        # (data.bit_length() + 7) // 8 计算字节数组的长度
        # byteorder='big' 指定字节顺序为大端序
        data = data.to_bytes((data.bit_length() + 7) // 8, byteorder='big')
    # 检查输入的数据类型是否为字符串
    elif isinstance(data, str):
        # 如果是字符串，将其编码为字节数组
        data = data.encode()

    # 使用 SHA256 算法对字节数组进行哈希
    hash_obj = SHA1.new(data)
    # 将哈希值的十六进制表示转换为整数并返回
    return int(hash_obj.hexdigest(), 16)


def H1(message):
    # return int.from_bytes(hashlib.sha1(message).digest(), byteorder='big')  # 将sha1生成的字节值转换为int
    return H(message)


# def H2(message):
#     return int.from_bytes(hashlib.sha256(message).digest(), byteorder='big')  # 将sha256生成的字节值转换为int,256的输出字节长度为256位

def H2(data):
    # 检查输入的数据类型是否为整数
    if isinstance(data, int):
        data = data
    # 检查输入的数据类型是否为字符串
    elif isinstance(data, str):
        # 如果是字符串，将其编码为字节数组
        data_byte = data.encode()
        data = byteToInt(data_byte)  # 字节转int
    elif isinstance(data, bytes):
        data = byteToInt(data)  # 字节转int

    hash_obj = data ^ (data >> 16) ^ (data >> 8)
    return hash_obj


def H3(message):
    # return int.from_bytes(hashlib.sha1(message).digest(), byteorder='big')  # 将sha1生成的字节值转换为int
    return H(message)


# 加密函数
def encrypt(key, message):
    plaintext = pickle.dumps(message)  # 将message的元组修改为字节类型
    cipher = AES.new(key, AES.MODE_GCM)  # 创建一个加密对象
    nonce = cipher.nonce  # 创建一次性数字，用于保证在密钥和明文相同的情况下，加密的结果不一致
    ciphertext, tag = cipher.encrypt_and_digest(plaintext)  # ciphertext是密文，tag是认证标签，保证密文的完整性和真实性
    return ciphertext, nonce, tag


# 解密函数
def decrypt(key, ciphertext, nonce, tag):
    cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
    t = cipher.decrypt(ciphertext)  # 解密，结果为字节
    plaintext = pickle.loads(t)  # 将解密后的明文修转化为元组
    try:
        cipher.verify(tag)  # 验证tag是否相同，相同则未被修改
        return plaintext[0], plaintext[1], plaintext[2]  # 将元组中的数据拿出，在此代码中为加密后的数据XB、YB、TB等值
    except ValueError:
        print("key is not valid")
        return False


# 时间戳验证
def timestamp_is_fresh(timestamp):
    now = time.time()
    if now - timestamp > 60:  # 60s内皆为新鲜
        return False
    else:
        return True


# 生成hw D
def generate_password_file(password, idd):
    s = get_random_bytes(1)  # 随机生成s
    # s_int = int.from_bytes(s, byteorder="big")  # 将随机生成的随机数从字节转换成int
    s_int = byteToInt(s)  # 将随机生成的随机数从字节转换成int
    hw = H1(password) ^ s_int
    D = H2(s) ^ byteToInt(idd)  # idd是字节，需要转换成int类型进行异或
    return hw, D


def session():
    rA = get_random_bytes(1)
    # rA_int = int.from_bytes(rA, byteorder="big")
    rA_int = byteToInt(rA)
    hwA, DA = generate_password_file(pw, idA)
    XA = H2(intToByte(hwA)) ^ H2(rA)
    YA = g ** rA_int
    YA_new = (g ** rA_int)*H2(pw)
    # c = H2(pw)
    TA = time.time()
    # XA、YA、TA生成，发送给B
    rB = get_random_bytes(1)
    rB_int = byteToInt(rB)  # 将字节转换为int方便后续计算
    hwB, DB = generate_password_file(pw, idB)  # 生成的值为int
    XB = H2(intToByte(hwB)) ^ H2(rB)  # 需要将hw的int类型转换为byte
    YB = g ** rB_int
    if not timestamp_is_fresh(TA):
        print("TA has expired")
    else:
        YA1 = YA_new//H2(pw)
        tkB = XA ^ DB ^ byteToInt(idB) ^ H2(rB) ^ H2(H3(YA1 ** rB_int))
        # 共20位，差4位 用0补齐
        tkB_byte = intToByte(tkB)
        # 补齐字节
        padding = (8 - len(tkB_byte) % 8) % 8  # 计算补齐的字节数，8的倍数
        tkB_byte += b'\x00' * padding  # 补齐后的值
        # print(f"tkB={tkB_byte}")
        # tkB = b't2Fb9ZvN8hsPQcRxg6oWnm45EDjKLYHa'
        TB = time.time()
        CB, nonce, tag = encrypt(tkB_byte, (TB, hwB, rB))
        # UB验证完TA之后将TB、hwB、rB加密，并将XB、YB、TB、CB给到UA
        if not timestamp_is_fresh(TB):
            print("TB has expired")
        else:
            tkA = XB ^ DA ^ byteToInt(idA) ^ H2(rA) ^ H2(H3(YB ** rA_int))
            tkA_byte = intToByte(tkA)
            # 补齐字节
            padding = (8 - len(tkA_byte) % 8) % 8
            tkA_byte += b'\x00' * padding
            # print(f"tkA={tkA_byte}")
            # tkA = b't2Fb9ZvN8hsPQcRxg6oWnm45EDjKLYHa'
            dec_TB, dec_hwB, dec_rB = decrypt(tkA_byte, CB, nonce, tag)
            if dec_TB != TB:
                print("dec_TB does not match TB, is not valid")
            else:
                if YB != g ** byteToInt(dec_rB) or XB != H2(intToByte(dec_hwB)) ^ H2(dec_rB):
                    print("YB does not match or XB does not match")
                else:
                    CA, nonce_CA, tag_CA = encrypt(tkA_byte, (TA, hwA, rA))
                    sidA = str(XA) + str(YA) + str(XB) + str(YB)  # XA、YA、XB、YB都是经过哈希或者幂运算的，都是int类型
                    KA = H(str(tkA).encode('utf-8') + sidA.encode('utf-8'))
                    # print('KA=', KA)
                    # UA 执行完毕，UB收到CA之后进行验证和计算
                    dec_TA, dec_hwA, dec_rA = decrypt(tkB_byte, CA, nonce_CA, tag_CA)
                    if dec_TA != TA:
                        print("dec_TA does not match, is not valid")
                    else:
                        if YA != g ** byteToInt(dec_rA) or XA != H2(intToByte(dec_hwA)) ^ H2(dec_rA):
                            print("YA does not match or XA does not match")
                        else:
                            sidB = str(XA) + str(YA) + str(XB) + str(YB)  # XA、YA、XB、YB都是经过哈希或者幂运算的，都是int类型
                            KB = H(str(tkB).encode('utf-8') + sidB.encode('utf-8'))
                            # print("KB=", KB)
                            return KA, KB


if __name__ == '__main__':
    # generate_password_file(pw, idA)
    # start_time = time.time()
    # KA, KB = session()
    # assert KA == KB, "Session keys do not match!"
    # print("Key exchange successful!")
    # end_time = time.time()
    # print("运行时间：", end_time - start_time)

    # 计算运行10000次平均的消耗的时间
    i = 0
    m = 100000
    sum_time = 0
    while (i < m):
        start_time = time.time()
        KA, KB = session()
        assert KA == KB, "Session keys do not match!"
        print("第" + str(i + 1) + "次Key exchange successful!")
        end_time = time.time()
        run_time = end_time - start_time
        print("第" + str(i + 1) + "运行时间：", run_time)
        sum_time += run_time
        i = i + 1
    print("运行" + str(m) + "次的时间：", sum_time)
    print("运行" + str(m) + "次的平均时间：", sum_time / m)
